from typing import TypedDict, Any

from langchain_core.messages import SystemMessage, BaseMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent

from src.ai.langchain.init_llm import get_llm


# 1. 定义工具
@tool
def get_weather(city: str) -> str:
    """返回指定城市的天气（中文城市名）"""
    return f"{city}是一个下雨天"


@tool
def get_temperatura(city: str) -> str:
    """返回指定城市的当前气温（中文城市名）"""
    return f"{city}的气温是25度"


# 2. 定义 prompt
prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessage(content="你是一个中文智能助手，严格使用用户提供的城市名，不要替换为其他城市。"),
        ("placeholder", "{messages}"),
    ]
)

# 3. 创建 agent（ReAct 风格）
llm = get_llm()
agent: Runnable = create_react_agent(
    model=llm,
    tools=[get_weather, get_temperatura],
    prompt=prompt,
)


# 4. 定义输入类型
class AgentInput(TypedDict):
    input: str


# 5. 封装调用函数
def run_agent(user_input: str) -> str:
    query: AgentInput = {"input": user_input}
    result: dict[str, Any] = agent.invoke(query)

    messages = result.get("messages", [])
    if not messages:
        return "无响应"

    last_message: BaseMessage = messages[-1]

    # 如果是 AIMessage（标准 LLM 输出）
    if isinstance(last_message, AIMessage):
        return last_message.content or "（无内容）"

    # 否则就打印字符串形式（例如 ToolMessage、FunctionMessage）
    return str(last_message)


response = run_agent("天津的天气如何？")
print(response)
